import torch as pt
import numpy as np

WEIGHT_PATH = '../../../../../large_data/model/darknet19/darknet19-deepBakSu-e1b3ec1e.pth'


class ConvBnRelu(pt.nn.Module):

    def __init__(self, in_ch, out_ch, ksize=3, stride=1, padding=1, bias=False, **kwargs):
        super().__init__(**kwargs)
        self.conv = pt.nn.Conv2d(in_ch, out_ch, ksize, stride, padding, bias=bias)
        self.bn = pt.nn.BatchNorm2d(out_ch)
        self.relu = pt.nn.LeakyReLU(inplace=True, negative_slope=0.1)

    def forward(self, inputs):
        x = self.conv(inputs)
        x = self.bn(x)
        x = self.relu(x)
        return x


cfg = [
    32, 'm',
    64, 'm',
    128, 64, 128, 'm',
    256, 128, 256, 'm',
    512, 256, 512, 256, 512, 'm',
    1024, 512, 1024, 512, 1024,
]


class MyDarknet19(pt.nn.Module):

    def __init__(self, in_ch, n_cls, is_pretrained, **kwargs):
        super().__init__(**kwargs)
        layers = []
        ksize = [1, 3]
        ksize_idx = True
        padding = [0, 1]
        padding_idx = True
        for cfg_v in cfg:
            if cfg_v == 'm':
                layers.append(pt.nn.MaxPool2d(2, 2, 0))
            else:
                layers.append(ConvBnRelu(in_ch, cfg_v, ksize[ksize_idx], 1, padding[padding_idx]))
                in_ch = cfg_v
            ksize_idx = not ksize_idx
            padding_idx = not padding_idx
        self.features = pt.nn.Sequential(*layers)
        self.clf_conv = pt.nn.Conv2d(cfg_v, n_cls, 1, 1, 0)
        self.clf_avg = pt.nn.AdaptiveAvgPool2d(1)
        self.clf_softmax = pt.nn.Softmax(dim=0)

        if is_pretrained:
            dict = {}
            keys = self.state_dict().keys()
            values = pt.load(WEIGHT_PATH).values()
            for k, v in zip(keys, values):
                dict[k] = v
            self.load_state_dict(dict)

    def forward(self, inputs):
        x = self.features(inputs)
        print('x before clf_conv', x.size())
        x = self.clf_conv(x)
        print('x after clf_conv', x.size())
        x = self.clf_avg(x)
        print('x after clf_avg', x.size())
        x = self.clf_softmax(x)
        print('x after clf_softmax', x.size())
        return x

if __name__ == '__main__':

    net = MyDarknet19(in_ch=3, n_cls=1000, is_pretrained=True)
    print(net)

    x = pt.zeros((4,3,224,224))
    out = net(x)
    print(out.size())
